import ujson as json
import ubinascii as binascii
from usr.qframe.threading import Lock, Condition


class SerialNumber(object):
    max_number = 0xFFFF
    lock = Lock()
    iterator = iter(range(max_number + 1))

    @classmethod
    def get(cls):
        with cls.lock:
            try:
                return next(cls.iterator)
            except StopIteration:
                cls.iterator = iter(range(cls.max_number + 1))
                return next(cls.iterator)


class Payload(object):

    def __init__(self, cmd, kwargs=None):
        self.cmd = cmd
        self.kwargs = kwargs or {}

    def to_bytes(self):
        return json.dumps({'cmd': self.cmd, 'kwargs': self.kwargs}).encode()

    @classmethod
    def from_bytes(cls, raw):
        payload = json.loads(raw.decode())
        return cls(
            cmd=payload['cmd'],
            kwargs=payload['kwargs']
        )


class Message(object):

    def __init__(self, cmd, kwargs=None, serial_number=None, ack=False):
        self.cmd = cmd
        self.kwargs = kwargs or {}
        self.payload = Payload(cmd, kwargs)
        self.serial_number = SerialNumber.get() if serial_number is None else serial_number
        self.ack = ack

    def __repr__(self):
        return '{}({}, kwargs={}, serial_number={}, ack={})'.format(
            type(self).__name__,
            self.cmd,
            self.kwargs,
            self.serial_number,
            self.ack
        )

    def to_bytes(self):
        head = 0x00
        if self.ack:
            head |= 0x01
        data = b'$'
        data += head.to_bytes(1, 'big')
        data += self.serial_number.to_bytes(2, 'big')
        body = self.payload.to_bytes()
        data += len(body).to_bytes(2, 'big') + body
        data += self.gen_crc(data).to_bytes(1, 'big')
        return b'[' + self.escape(data) + b']'

    @staticmethod
    def escape(data):
        origin = b''
        for one in data:
            if one == 0x25:
                origin += b'\x25\x00'
            elif one == 0x5B:
                origin += b'\x25\x01'
            elif one == 0x5D:
                origin += b'\x25\x02'
            else:
                origin += one.to_bytes(1, 'big')
        return origin

    @staticmethod
    def revert(data):
        origin = b''
        index = 0
        while index <= len(data) - 1:
            one = data[index]
            if one == 0x25:
                if data[index + 1] == 0x00:
                    origin += b'\x25'
                elif data[index + 1] == 0x01:
                    origin += b'\x5B'
                elif data[index + 1] == 0x02:
                    origin += b'\x5D'
                else:
                    raise ValueError('revert error at {} bytes, 0x01 or 0x02 should be followed.'.format(index))
                index += 2
            else:
                origin += one.to_bytes(1, 'big')
                index += 1
        return origin

    @classmethod
    def from_bytes(cls, raw):
        if len(raw) < 10:
            raise ValueError('message less than 10 bytes.')
        if raw[0] != 0x5B or raw[-1] != 0x5D:
            raise ValueError('boundary error, not a valid message.')
        data = cls.revert(raw[1:-1])
        crc = data[-1]
        if crc != cls.gen_crc(data[:-1]):
            raise ValueError('CRC validate FAILED.')
        if data[0] != 0x24:
            raise ValueError('identifier error, not a valid message.')
        is_ack = bool(data[1] & 0x01)
        serial_number = int.from_bytes(data[2:4], 'big')
        payload_length = int.from_bytes(data[4:6], 'big')
        payload = Payload.from_bytes(data[6:6+payload_length])
        self = cls(
            payload.cmd,
            kwargs=payload.kwargs,
            serial_number=serial_number,
            ack=is_ack
        )
        return self

    @staticmethod
    def gen_crc(data):
        crc = data[0]
        for one in data[1:]:
            crc ^= one
        return crc

    def to_hex(self):
        return binascii.hexlify(self.to_bytes()).decode().upper()

    @classmethod
    def from_hex(cls, string):
        return cls.from_bytes(binascii.unhexlify(string))

    @classmethod
    def make_request(cls, cmd, **kwargs):
        return cls(cmd, kwargs)

    def make_response(self, **kwargs):
        if self.ack:
            raise ValueError('ack message can not make response.')
        return type(self)(
            self.cmd,
            kwargs=kwargs,
            serial_number=self.serial_number,
            ack=True
        )


class Parser(object):

    def __init__(self, load=True):
        self.__buffer = b''
        self.__message_list = []
        self.__load = load

    def __handle_message(self, raw):
        if self.__load:
            try:
                self.__message_list.append(Message.from_bytes(raw))
            except Exception as e:
                print('protocol message parse error: {}; we just pass it!'.format(e))
        else:
            self.__message_list.append(raw)

    def parse(self, data):
        self.__buffer += data
        temp = b''
        for c in self.__buffer:
            if c == 0x5B:
                if 0x5B in temp:
                    temp = b'['
                    continue
                temp += b'['
            elif c == 0x5D:
                if 0x5B not in temp:
                    temp = b''
                    continue
                temp += b']'
                self.__handle_message(temp)
                temp = b''
            else:
                if b'[' not in temp:
                    continue
                temp += c.to_bytes(1, 'big')
        self.__buffer = temp

    @property
    def messages(self):
        rv, self.__message_list = self.__message_list, []
        return rv

    def clear(self):
        self.__buffer = b''


class RespHelper(Condition):

    def __init__(self):
        self.__ack_items = {}
        super().__init__()

    def get(self, request, timeout=None):
        """accept a request and return response matched or none"""
        self.__ack_items[request] = None
        self.wait_for(lambda: self.__ack_items[request] is not None, timeout=timeout)
        return self.__ack_items.pop(request)

    def put(self, response):
        """accept a response and match it with request if possible"""
        for request in self.__ack_items.keys():
            if self.validate(request, response):
                self.__ack_items[request] = response
                self.notify_all()
                break

    @staticmethod
    def validate(request, response):
        return request.serial_number == response.serial_number and request.cmd == response.cmd


if __name__ == '__main__':
    m = Message.make_request('TEST', a=1, b=2)
    print(m)
